package cn.tannn.tgateway.plugin;

import cn.tannn.tgateway.AbstractGatewayPlugin;
import cn.tannn.tgateway.GatewayPluginChain;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * direct proxy plugin (直接代理
 *
 * @author <a href="https://t.tannn.cn/">tan</a>
 * @date 2024/5/30 下午4:21
 */
@Component("direct")
@Slf4j
public class DirectPlugin extends AbstractGatewayPlugin {


    /**
     * plugin name
     */
    public static final String NAME = "direct";
    /**
     * request prefix
     */
    private String prefix = GATEWAY_PREFIX + "/" + NAME + "/";


    @Override
    public boolean doSupport(ServerWebExchange exchange) {
        return exchange.getRequest().getPath().value().startsWith(prefix);
    }

    @Override
    public Mono<Void> doHandle(ServerWebExchange exchange, GatewayPluginChain chain) {
        log.debug("=======>>>>>>> [DirectPlugin] ...");
        // 后端地址
        String backend = exchange.getRequest().getQueryParams().getFirst("backend");


        exchange.getResponse().getHeaders().add("Content-Type", "application/json");
        exchange.getResponse().getHeaders().add("tan.gw.version", "v1.0.0");
        exchange.getResponse().getHeaders().add("tan.gw.plugin", getName());

        if (backend == null || backend.isEmpty()) {
            Flux<DataBuffer> requestBody = exchange.getRequest().getBody();
            return requestBody.flatMap(x -> exchange.getResponse().writeWith(Mono.just(x)))
                    // 继续执行 chain 逻辑，不打断 plugin链
                    .then(chain.handle(exchange));
        }

        HttpMethod httpMethod = exchange.getRequest().getMethod();
        WebClient client = WebClient.create();

        ;

        WebClient.RequestHeadersSpec<?> spec = switch (httpMethod.name()) {
            case "GET" -> client.get()
                    .uri(backend, uriBuilder -> uriBuilder
                            .queryParams(exchange.getRequest().getQueryParams())
                            .build());
            case "POST" -> client.post().uri(backend);
            case "PUT" -> client.put().uri(backend)
                    .body(exchange.getRequest().getBody(), DataBuffer.class);
            case "DELETE" -> client.delete()
                    .uri(backend, uriBuilder -> uriBuilder
                            .queryParams(exchange.getRequest().getQueryParams())
                            .build());

            default -> throw new IllegalArgumentException("Unsupported HTTP method: " + httpMethod);
        };


        Mono<ResponseEntity<String>> entity = spec.header("Content-Type", "application/json")
                .retrieve().toEntity(String.class);

        Mono<String> body = entity.mapNotNull(ResponseEntity::getBody);
        return body.flatMap(x -> exchange.getResponse()
                        .writeWith(Mono.just(exchange.getResponse().bufferFactory().wrap(x.getBytes()))))
                .then(chain.handle(exchange));
    }

    @Override
    public String getName() {
        return NAME;
    }
}
